﻿//
// Created by QU on 24-10-6.
//

#include "server/server/WorkerPool.h"
#include "server/server/Worker.h"
#include "server/server/Session.h"

#include <iostream>
#include <numeric>
#include <algorithm>
#include <thread>
#include <future>
#include <chrono>
#include <utility>

WorkerPool::WorkerPool()
    : m_acceptor_io_work_guard(asio::make_work_guard(m_acceptor_io_context))
      , m_shared_worker_stop_source(std::make_shared<std::stop_source>())
{
}

WorkerPool::~WorkerPool() = default;

bool WorkerPool::startAcceptor(const uint16_t port)
{
    try
    {
        m_acceptor = std::make_unique<tcp::acceptor>(m_acceptor_io_context, tcp::endpoint(tcp::v4(), port));

        m_acceptor_thread = std::jthread(
            [this](std::stop_token stoken)
            {
                doAccept(stoken);
                m_acceptor_io_context.run();
                std::cout << "Acceptor io_context exited." << std::endl;
            });

        return true;
    }
    catch (const std::exception& e)
    {
        std::cerr << "Error starting acceptor on port " << port << ": " << e.what() << std::endl;
        return false;
    }
}

void WorkerPool::doAccept(std::stop_token& stoken)
{
    if (stoken.stop_requested())
    {
        return;
    }

    m_acceptor->async_accept(
        [this, stoken](const boost::system::error_code& error, tcp::socket socket) mutable
        {
            if (stoken.stop_requested())
            {
                std::cout << "Stop requested, cancelling accept operation." << std::endl;
                return;
            }

            if (!error)
            {
                handleAccept(error, std::move(socket));
            }
            else if (error == asio::error::operation_aborted)
            {
                std::cout << "WorkerPool::doAccept(): callback: Connection closed: operation_aborted." << std::endl;
                return;
            }
            else
            {
                std::cerr << "Accept error: " << error.message() << std::endl;
            }

            if (!stoken.stop_requested())
            {
                doAccept(stoken); // Continue accepting connections if stop is not requested
            }
        });
}


void WorkerPool::handleAccept(const boost::system::error_code& error, tcp::socket&& socket)
{
    if (error)
    {
        std::cerr << "Error handling accept: " << error.message() << std::endl;
        return;
    }

    try
    {
        printRemoteSockInformation(socket);
        Worker& worker = getLeastLoadedWorker();
        std::cout << "[debug] Worker Session Count: " <<worker.getSessionCount() << std::endl;

        const auto session = std::make_shared<Session>(std::move(socket), worker,
                                                       worker.getSharedSessionStopSource()->get_token());
        worker.addSession(session->getUuid(), session);
        session->startSession();
    }
    catch (const std::exception& e)
    {
        std::cerr << "Exception in handleAccept: " << e.what() << std::endl;
    }
}

bool WorkerPool::stopAcceptor()
{
    try
    {
        std::cout << "Requesting stop for acceptor..." << std::endl;
        m_acceptor_thread.request_stop();

        if (m_acceptor)
        {
            std::cout << "Cancelling acceptor..." << std::endl;
            boost::system::error_code ec;
            m_acceptor->cancel(ec); // Cancels any asynchronous operations.      NOLINT(*-unused-return-value)
            if (ec) std::cerr << "Error cancelling acceptor: " << ec.message() << std::endl;

            std::cout << "Closing acceptor..." << std::endl;
            m_acceptor->close(ec); // Close the acceptor.      NOLINT(*-unused-return-value)
            if (ec) std::cerr << "Error closing acceptor: " << ec.message() << std::endl;
        }

        std::cout << "Stopping io_context..." << std::endl;
        m_acceptor_io_context.stop(); // Explicitly stop the io_context

        m_acceptor_io_work_guard.reset(); // Reset work guard after stopping io_context
        std::cout << "Work guard reset." << std::endl;

        if (m_acceptor_thread.joinable())
        {
            std::cout << "Waiting for acceptor thread to stop..." << std::endl;
            m_acceptor_thread.join(); // Wait for the thread to finish
        }

        std::cout << "Acceptor stopped successfully." << std::endl;
        return true;
    }
    catch (const std::exception& e)
    {
        std::cerr << "Error stopping acceptor: " << e.what() << std::endl;
        return false;
    }
}


bool WorkerPool::startWorkers(const size_t worker_thread_count)
{
    try
    {
        for (size_t i = 0; i < worker_thread_count; ++i)
        {
            auto io_context = std::make_unique<asio::io_context>();
            auto worker = std::make_unique<Worker>(*io_context, m_shared_worker_stop_source->get_token());
            worker->start();
            {
                std::scoped_lock lock(m_mutex);
                m_worker_io_context_s.push_back(std::move(io_context));
                m_worker_s.push_back(std::move(worker));
            }
        }
        std::cout << "Started " << worker_thread_count << " workers." << std::endl;
        return true;
    }
    catch (const std::exception& e)
    {
        std::cerr << "Error starting workers: " << e.what() << std::endl;
        return false;
    }
}

bool WorkerPool::requestStopWorkers()
{
    try
    {
        std::cout << "Requesting stop workers." << std::endl;
        m_shared_worker_stop_source->request_stop();

        const auto start_time = std::chrono::steady_clock::now();
        constexpr auto timeout = std::chrono::seconds(30); // 30秒超时

        for (size_t i = 0; i < m_worker_s.size(); ++i)
        {
            while (m_worker_s[i]->isWorking())
            {
                if (std::chrono::steady_clock::now() - start_time > timeout)
                {
                    std::cout << "Worker " << i << " did not stop in time. Forcing stop." << std::endl;
                    m_worker_s[i]->forceStop();
                    break;
                }
                std::this_thread::sleep_for(std::chrono::milliseconds(100));
            }
            std::cout << "Worker " << i << " stopped." << std::endl;
        }

        std::cout << "All workers stopped." << std::endl;

        // 清理工作线程和IO上下文
        {
            std::scoped_lock lock(m_mutex);
            m_worker_s.clear();
            m_worker_io_context_s.clear();
        }

        std::cout << "Cleared worker containers." << std::endl;

        return true;
    }
    catch (const std::exception& e)
    {
        std::cerr << "Error requesting stop for workers: " << e.what() << std::endl;
        return false;
    }
}


bool WorkerPool::forceStopWorkers()
{
    try
    {
        m_shared_worker_stop_source->request_stop();

        for (const auto& worker : m_worker_s)
        {
            worker->forceStop();
        }

        {
            std::scoped_lock lock(m_mutex);
            m_worker_s.clear();
            m_worker_io_context_s.clear();
        }

        std::cout << "Forcefully stopped all workers." << std::endl;
        return true;
    }
    catch (const std::exception& e)
    {
        std::cerr << "Error force stopping workers: " << e.what() << std::endl;
        return false;
    }
}


bool WorkerPool::startAll(const uint16_t port, size_t worker_count)
{
    if (worker_count == 0)
    {
        worker_count = std::thread::hardware_concurrency();
        std::cout << "[INFO] Using default thread count: " << worker_count << std::endl;
    }

    if (!startWorkers(worker_count))
    {
        return false;
    }

    if (!startAcceptor(port))
    {
        return false;
    }

    return true;
}

void WorkerPool::requestStopAll()
{
    std::cout << "Stopping worker pool..." << std::endl;
    stopAcceptor();
    requestStopWorkers();
}

void WorkerPool::forceStopAll()
{
    std::cout << "Force Stopping worker pool..." << std::endl;
    stopAcceptor();
    forceStopWorkers();
}

asio::io_context& WorkerPool::getAcceptorIOContext()
{
    return m_acceptor_io_context;
}


void WorkerPool::printRemoteSockInformation(const ip::tcp::socket& sock)
{
    try
    {
        const auto remote_endpoint = sock.remote_endpoint();
        std::cout << "New connection from: " << remote_endpoint.address().to_string()
            << ":" << remote_endpoint.port() << std::endl;
    }
    catch (const std::exception& e)
    {
        std::cerr << "Error getting remote endpoint: " << e.what() << std::endl;
    }
}

size_t WorkerPool::getActiveWorkerCount() const
{
    std::scoped_lock lock(m_mutex);
    return m_worker_s.size();
}

std::vector<size_t> WorkerPool::getWorkerSessionCounts() const
{
    std::scoped_lock lock(m_mutex);
    std::vector<size_t> counts;
    counts.reserve(m_worker_s.size());
    for (const auto& worker : m_worker_s)
    {
        counts.push_back(worker->getSessionCount());
    }
    return counts;
}

Worker& WorkerPool::getLeastLoadedWorker()
{
    std::scoped_lock lock(m_mutex);
    const auto it = std::ranges::min_element(m_worker_s,
                                             [](const auto& a, const auto& b)
                                             {
                                                 return a->getSessionCount() < b->getSessionCount();
                                             });

    if (it != m_worker_s.end())
    {
        return **it;
    }

    throw std::runtime_error("No workers available");
}

size_t WorkerPool::getTotalSessionCount() const
{
    std::scoped_lock lock(m_mutex);
    return std::accumulate(m_worker_s.begin(), m_worker_s.end(), static_cast<size_t>(0),
                           [](size_t sum, const auto& worker)
                           {
                               return sum + worker->getSessionCount();
                           });
}
